package top.codetheory._05_二叉搜索树;

import java.util.Comparator;

/**
 * @author:
 * @since: 2023/12/14 14:43
 **/
public class AVLTree<E> extends BBST<E>{
    public AVLTree() {
        this(null);
    }
    public AVLTree(Comparator<E> comparator) {
        super(comparator);
    }

    /**
     * 更新之后做的事, 保持平衡
     * @param node 新添加的节点
     */
    @Override
    protected void afterAdd(Node<E> node) {
        while ((node = node.parent) != null) {  // 直接用父节点覆盖当前节点
            if (isBalanced(node)) {  // 判断当前节点是否平衡
                updateHeight(node);  // 是平衡的, 那么就更新一下高度
            } else { // 不是平衡的, 那么需要恢复高度
//                reBalance(node);  // 恢复平衡, 来到这个位置的节点应该的爷爷节点, parent, 父节点原来肯定是平衡的, 再加一个节点, 父节点也是平衡的, grandfather 可能不平衡
                reBalanceUnification(node);  // 统一的恢复平衡方法
                break; // 爷爷恢复平衡了, 那么整棵树就恢复平衡了
            }

        }
    }

    @Override
    protected void afterRemove(Node<E> node, Node<E> replacement) {
        while ((node = node.parent) != null) {  // 直接用父节点覆盖当前节点
            if (isBalanced(node)) {  // 判断当前节点是否平衡
                updateHeight(node);  // 是平衡的, 那么就更新一下高度
            } else { // 不是平衡的, 那么需要恢复高度
//                reBalance(node);  // 恢复平衡, 来到这个位置的节点应该的爷爷节点, parent, 父节点原来肯定是平衡的, 再加一个节点, 父节点也是平衡的, grandfather 可能不平衡
                reBalanceUnification(node);  // 统一的恢复平衡方法
                // 去掉了break, 因为这里祖先节点也可能失衡
            }
        }
    }

    /**
     * 统一恢复平衡的方法
     * @param node grandfather节点 高度最低的那个不平衡的节点
     */
    protected void reBalanceUnification(Node<E> node) {
        AVLNode<E> grand = (AVLNode<E>) node;
        AVLNode<E> parent = grand.tallerChild();
        AVLNode<E> child = parent.tallerChild();
        if (parent.isLeftChildren()) {  // L
            if (child.isLeftChildren()) {  // LL
                rotate(grand, child.left, child, child.right, parent, parent.right, grand, grand.right);
            } else {  // LR
                rotate(grand, parent.left, parent, child.left, child, child.right, grand, grand.right);
            }
        } else { // R
            if (child.isLeftChildren()) {  // RL
                rotate(grand, grand.left, grand, child.left, child, child.right, parent, parent.right);
            } else {  // RR
                rotate(grand, grand.left, grand, parent.left, parent, child.left, child, child.right);
            }
        }
    }

    /**
     * 恢复平衡
     * @param node grandfather节点 高度最低的那个不平衡的节点
     */
    private void reBalance(Node<E> node) {
        AVLNode<E> grand = (AVLNode<E>) node;
        AVLNode<E> parent = grand.tallerChild();
        AVLNode<E> child = parent.tallerChild();
        if (parent.isLeftChildren()) {  // L
            if (child.isLeftChildren()) {  // LL
                rotateRight(grand);
            } else {  // LR
                rotateLeft(parent);
                rotateRight(grand);
            }
        } else { // R
            if (child.isLeftChildren()) {  // RL
                rotateRight(parent);
                rotateLeft(grand);
            } else {  // RR
                rotateLeft(grand);
            }
        }
    }



    /**
     * 统一旋转
     * @param r 原来子树的根节点
     * @param a b的左
     * @param b d的左子树, 根节点
     * @param c b的右
     * @param d 旋转后的根节点
     * @param e f的左
     * @param f d的右子树, 根节点
     * @param g f的右
     */
    protected void rotate(
            Node<E> r,
            Node<E> a, Node<E> b, Node<E> c,
            Node<E> d,
            Node<E> e, Node<E> f, Node<E> g
    ) {
        super.rotate(r, a, b, c, d, e, f, g);
        // b 的左右子树发生了变化, 所以要更新一下b的高度
        updateHeight(b);
        updateHeight(f);
        updateHeight(d);
    }


    /**
     * 旋转之后的操作
     * @param grand 爷爷节点
     * @param parent 父节点
     * @param node n节点
     */
    @Override
    protected void afterRotate(Node<E> grand, Node<E> parent, Node<E> node) {
        super.afterRotate(grand, parent, node);
        // 更新g 和 p的高度
        updateHeight(grand);
        updateHeight(parent);
    }

    /**
     * 更新节点的高度
     * @param node 需要更新的节点
     */
    protected void updateHeight(Node<E> node) {
        ((AVLTree.AVLNode<E>)node).updateHeight();
    }

    /**
     * 判断节点是否平衡
     * @param node 需要判断的节点
     * @return 是否平衡
     */
    private boolean isBalanced(Node<E> node) {
        return Math.abs(((AVLNode<E>)node).balanceFactor()) <= 1;
    }

    @Override
    protected Node<E> createNode(E element, Node<E> parent) {
        return new AVLNode<>(element, parent);
    }

    private static class AVLNode<E> extends Node<E> {
        int height = 1;  // 默认高度为1

        public AVLNode(E element, Node<E> parent) {
            super(element, parent);
        }

        /**
         * 计算平衡因子
         * @return 平衡因子
         */
        public int balanceFactor() {
            int leftHeight = left == null ? 0 : ((AVLNode<E>)left).height;
            int rightHeight = right == null ? 0 : ((AVLNode<E>)right).height;
            return leftHeight - rightHeight;
        }

        /**
         * 更新高度此节点的高度, 左子树与右子树的最大值
         */
        public void updateHeight() {
            int leftHeight = left == null ? 0 : ((AVLNode<E>)left).height;
            int rightHeight = right == null ? 0 : ((AVLNode<E>)right).height;
            height = 1 + Math.max(leftHeight, rightHeight);
        }

        /**
         * 获取该节点的最长的子节点
         *    1
         *  /  \
         * 2    3
         *       \
         *        4
         * @return 获取1的最长子节点 -> 3
         */
        public AVLNode<E> tallerChild() {
            int leftHeight = left == null ? 0 : ((AVLNode<E>)left).height;
            int rightHeight = right == null ? 0 : ((AVLNode<E>)right).height;
            return (AVLNode<E>) (leftHeight > rightHeight ? this.left : this.right);
        }
    }

    @Override
    public Object string(Object node) {
        AVLNode<E> myNode = (AVLNode<E>)node;
        String parentString = "null";
        if (myNode.parent != null) {
            parentString = myNode.parent.element.toString();
        }
        return myNode.element + "_p(" + parentString + ")" + "_h(" + myNode.height + ")";
    }
}